/*
 * SPDX-FileCopyrightText: 2022-2024 Andrew Gunnerson
 * SPDX-License-Identifier: GPL-3.0-only
 */

@file:OptIn(ExperimentalUnsignedTypes::class)

package com.sanji.audio_test.format

import android.media.MediaCodec
import android.media.MediaFormat
import android.system.Os
import android.system.OsConstants
import android.util.Log
import com.sanji.audio_test.writeFully
import java.io.FileDescriptor
import java.io.IOException
import java.nio.ByteBuffer

/**
 * 虚拟FLAC容器包装器，在编码完成时更新STREAMINFO中的持续时间字段
 *
 * [MediaCodec] 已经能生成格式良好的FLAC文件，因此该类直接将这些样本写入输出文件
 *
 * @param fd 输出文件描述符。该类不拥有该文件描述符的所有权，在调用 [stop] 并返回之前，
 * 不应在类外部操作该文件描述符
 */
class FlacContainer(private val fd: FileDescriptor) : Container {
    private var isStarted = false  // 容器是否已启动
    private var lastPresentationTimeUs = -1L  // 最后一个样本的显示时间戳（微秒）
    private var track = -1  // 轨道索引

    /**
     * 启动容器
     * 
     * @throws IllegalStateException 如果容器已经启动
     */
    override fun start() {
        if (isStarted) {
            throw IllegalStateException("Container already started")
        }

        // 将文件指针移到文件开始位置并清空文件
        Os.lseek(fd, 0, OsConstants.SEEK_SET)
        Os.ftruncate(fd, 0)

        isStarted = true
    }

    /**
     * 停止容器并更新FLAC头部信息
     * 
     * @throws IllegalStateException 如果容器未启动
     */
    override fun stop() {
        if (!isStarted) {
            throw IllegalStateException("Container not started")
        }

        isStarted = false

        // 如果有有效的最后时间戳，则更新头部的持续时间信息
        if (lastPresentationTimeUs >= 0) {
            Log.d(TAG, "设置头部中的持续时间字段")
            setHeaderDuration()
        }
    }

    /**
     * 释放容器资源
     */
    override fun release() {
        if (isStarted) {
            stop()
        }
    }

    /**
     * 添加媒体轨道
     * 
     * @param mediaFormat 媒体格式信息
     * @return 轨道索引（始终为0）
     * @throws IllegalStateException 如果容器已启动或轨道已添加
     */
    override fun addTrack(mediaFormat: MediaFormat): Int {
        if (isStarted) {
            throw IllegalStateException("Container already started")
        } else if (track >= 0) {
            throw IllegalStateException("Track already added")
        }

        track = 0
        return track
    }

    /**
     * 写入音频样本数据
     * 
     * @param trackIndex 轨道索引
     * @param byteBuffer 包含音频数据的字节缓冲区
     * @param bufferInfo 音频缓冲区信息
     * @throws IllegalStateException 如果容器未启动、未添加轨道或轨道索引无效
     */
    override fun writeSamples(trackIndex: Int, byteBuffer: ByteBuffer,
                              bufferInfo: MediaCodec.BufferInfo) {
        if (!isStarted) {
            throw IllegalStateException("Container not started")
        } else if (track < 0) {
            throw IllegalStateException("No track has been added")
        } else if (track != trackIndex) {
            throw IllegalStateException("Invalid track: $trackIndex")
        }

        // 将音频数据写入文件
        writeFully(fd, byteBuffer)

        // 检查是否为流的结尾
        if ((bufferInfo.flags and MediaCodec.BUFFER_FLAG_END_OF_STREAM) != 0) {
            lastPresentationTimeUs = bufferInfo.presentationTimeUs
            Log.d(TAG, "接收到EOF；最终显示时间戳: $lastPresentationTimeUs")
        }
    }

    /**
     * 将帧数写入FLAC文件的STREAMINFO元数据块
     *
     * @throws IOException 如果FLAC元数据无效或根据 [lastPresentationTimeUs] 计算出的帧数
     * 超出了36位整数的范围
     */
    private fun setHeaderDuration() {
        // 将文件指针移到文件开始位置
        Os.lseek(fd, 0, OsConstants.SEEK_SET)

        // 魔数(4字节)
        // + 元数据块头部(4字节)
        // + streaminfo块(34字节)
        val buf = UByteArray(42)

        // 读取FLAC头部信息
        if (Os.read(fd, buf.asByteArray(), 0, buf.size) != buf.size) {
            throw IOException("读取FLAC头部时遇到EOF")
        }

        // 验证魔数
        if (ByteBuffer.wrap(buf.asByteArray(), 0, 4) !=
            ByteBuffer.wrap(FLAC_MAGIC.asByteArray())) {
            throw IOException("未找到FLAC魔数")
        }

        // 验证第一个元数据块是STREAMINFO且大小正确
        if (buf[4] and 0x7fu != 0.toUByte()) {
            throw IOException("第一个元数据块不是STREAMINFO")
        }

        val streamInfoSize = buf[5].toUInt().shl(16) or
                buf[6].toUInt().shl(8) or buf[7].toUInt()
        if (streamInfoSize < 34u) {
            throw IOException("STREAMINFO块太小")
        }

        // 采样率字段是第18字节开始的20位整数
        val sampleRate = buf[18].toUInt().shl(12) or
                buf[19].toUInt().shl(4) or
                buf[20].toUInt().shr(4)

        // 这会略微低估持续时间，因为它没有考虑最后一次写入的持续时间
        val frames = lastPresentationTimeUs.toULong() * sampleRate / 1_000_000uL

        if (frames >= 2uL.shl(36)) {
            throw IOException("帧数无法在FLAC中表示: $frames")
        }

        // 总样本数字段是第21字节开始的36位整数（从第21字节的第4位开始）
        buf[21] = (buf[21] and 0xf0u) or (frames.shr(32) and 0xfu).toUByte()
        buf[22] = (frames.shr(24) and 0xffu).toUByte()
        buf[23] = (frames.shr(16) and 0xffu).toUByte()
        buf[24] = (frames.shr(8) and 0xffu).toUByte()
        buf[25] = (frames and 0xffu).toUByte()

        // 将更新后的帧数信息写回文件
        Os.lseek(fd, 21, OsConstants.SEEK_SET)
        writeFully(fd, buf.asByteArray(), 21, 5)
    }

    companion object {
        private val TAG = FlacContainer::class.java.simpleName
        private val FLAC_MAGIC = ubyteArrayOf(0x66u, 0x4cu, 0x61u, 0x43u) // fLaC
    }
}